package net.xt.testthread.config;

import lombok.Builder;
import lombok.Data;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadFactory;
import java.util.stream.Collectors;

/**
 * 线程池管理器
 */
public class ThreadPoolManager {

    private static Map<String, ThreadPoolTaskExecutor> threadPoolMap = new ConcurrentHashMap<>(16);
    //线程池核心线程数量（默认值）
    private static final int CORE_POOL_SIZE_DEFAULT = Runtime.getRuntime().availableProcessors() * 2;
    //线程池最大线程数量（默认值）
    private static final int MAX_POOL_SIZE_DEFAULT = CORE_POOL_SIZE_DEFAULT * 2;
    //线程池默认队列大小1000（默认值）
    private static final int QUEUE_CAPACITY_DEFAULT = 1000;
    //线程池中的线程默认存活时间5分钟（默认值）
    private static final int KEEP_ALIVE_SECONDS_DEFAULT = 300;

    public static ThreadPoolTaskExecutor newThreadPool(String name) {
        return newThreadPool(name, CORE_POOL_SIZE_DEFAULT, MAX_POOL_SIZE_DEFAULT);
    }

    public static ThreadPoolTaskExecutor newThreadPool(String name, int corePoolSize) {
        return newThreadPool(name, corePoolSize, corePoolSize);
    }

    public static ThreadPoolTaskExecutor newThreadPool(String name, int corePoolSize, int maxPoolSize) {
        return newThreadPool(name, corePoolSize, maxPoolSize, QUEUE_CAPACITY_DEFAULT, KEEP_ALIVE_SECONDS_DEFAULT,
                null, null);
    }

    /**
     * @param name                     线程池名称
     * @param corePoolSize             核心线程数d
     * @param maxPoolSize              最大线程数
     * @param queueCapacity            队列大小
     * @param keepAliveSeconds         线程池存活时间（秒）
     * @param threadFactory            线程工厂
     * @param rejectedExecutionHandler 拒绝策略
     */
    public static ThreadPoolTaskExecutor newThreadPool(String name, int corePoolSize,
                                                       int maxPoolSize, int queueCapacity, int keepAliveSeconds,
                                                       ThreadFactory threadFactory, RejectedExecutionHandler rejectedExecutionHandler) {
        return threadPoolMap.computeIfAbsent(name, threadGroupName -> {
            ThreadPoolTaskExecutor threadPoolExecutor = new ThreadPoolTaskExecutor();
            threadPoolExecutor.setCorePoolSize(corePoolSize);
            threadPoolExecutor.setMaxPoolSize(maxPoolSize);
            threadPoolExecutor.setQueueCapacity(queueCapacity);
            threadPoolExecutor.setKeepAliveSeconds(keepAliveSeconds);
            threadPoolExecutor.setThreadGroupName(name);
            threadPoolExecutor.setThreadFactory(threadFactory);
            threadPoolExecutor.afterPropertiesSet();
            threadPoolExecutor.setRejectedExecutionHandler(rejectedExecutionHandler);
            return threadPoolExecutor;
        });
    }

    /**
     * 获取所有线程池信息
     */
    public static List<ThreadPoolInfo> threadPoolInfoList() {
        return threadPoolMap.entrySet().stream()
                .map(entry -> {
                    ThreadPoolTaskExecutor poolTaskExecutor = threadPoolMap.get(entry.getKey());
                    return ThreadPoolInfo.builder()
                            .name(poolTaskExecutor.getThreadGroup().getName())
                            .activeCount(poolTaskExecutor.getActiveCount())
                            .maxPoolSize(poolTaskExecutor.getMaxPoolSize())
                            .corePoolSize(poolTaskExecutor.getCorePoolSize())
                            .queueCapacity(poolTaskExecutor.getQueueCapacity())
                            .queueSize(poolTaskExecutor.getQueueSize())
                            .build();
                })
                .collect(Collectors.toList());
    }

}